# -*- coding: utf-8 -*-
"""
项目名称：cotton_test
文件名称：app.py
创建时间：2025-06-08 20:47:50

系统用户：wyl
作　　者：無以菱
联系邮箱：huangjing2001.guet@qq.com
功能描述：棉花叶片疾病识别Web应用
- 本模块实现了一个基于Flask的简单Web应用，用于棉花叶片疾病识别测试。
- 应用允许用户上传棉花叶片图片，并使用预训练的卷积神经网络模型进行健康/患病分类。
"""


import os
import sys
import torch
import torch.nn as nn
from PIL import Image
import torchvision.transforms as transforms
from flask import Flask, request, render_template, jsonify, url_for
from werkzeug.utils import secure_filename  # 用于安全地处理文件名

# 添加项目根目录到系统路径，确保可以导入其他模块
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 导入模型和工具
from model import CustomCNN  # 导入自定义CNN模型
from utils.transform import transform  # 导入图像预处理转换

# 初始化Flask应用
app = Flask(__name__, static_url_path='/static')
# 配置上传文件夹路径
app.config['UPLOAD_FOLDER'] = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'uploads')
# 设置允许上传的文件类型
app.config['ALLOWED_EXTENSIONS'] = {'png', 'jpg', 'jpeg'}
# 限制上传文件大小为16MB
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024

# 确保上传目录存在
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

def allowed_file(filename):
    """
    检查文件扩展名是否在允许的列表中
    
    参数:
        filename (str): 要检查的文件名
        
    返回:
        bool: 如果文件扩展名允许则返回True，否则返回False
    """
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']

def load_model():
    """
    加载预训练的CNN模型
    
    函数会尝试从model目录加载模型，如果失败则尝试从根目录加载。
    
    返回:
        tuple: (model, device) - 加载的模型和计算设备
    """
    # 检查是否有可用的GPU，如果有则使用GPU，否则使用CPU
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(f"使用设备: {device}")
    
    # 创建模型实例 - 2个类别对应健康和患病
    model = CustomCNN(num_classes=2)
    
    # 加载预训练权重
    try:
        # 尝试从model目录加载
        model_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'model', 'cnn_model.pth')
        model.load_state_dict(torch.load(model_path, map_location=device))
        print(f"成功从{model_path}加载模型")
    except FileNotFoundError:
        # 如果找不到，尝试从根目录加载
        model_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'cnn_model.pth')
        model.load_state_dict(torch.load(model_path, map_location=device))
        print(f"成功从{model_path}加载模型")
    
    # 将模型设置为评估模式（不进行梯度计算和参数更新）
    model.eval()
    # 将模型移动到指定设备（GPU或CPU）
    model = model.to(device)
    
    return model, device

# 类别映射字典 - 将模型输出的数字标签映射为可读的类别名称
class_names = {
    0: "健康棉花叶片",
    1: "患病棉花叶片"
}

# 加载模型（应用启动时执行一次）
model, device = load_model()

@app.route('/')
def index():
    """
    渲染主页
    
    返回:
        str: 渲染后的HTML页面
    """
    return render_template('index.html')

@app.route('/predict', methods=['POST'])
def predict():
    """
    处理图片上传和预测请求
    
    接收上传的图片，使用模型进行预测，并返回预测结果
    
    返回:
        json: 包含预测结果的JSON响应
    """
    # 检查请求中是否包含文件
    if 'file' not in request.files:
        return jsonify({'error': '没有文件'})
    
    file = request.files['file']
    
    # 检查文件名是否为空
    if file.filename == '':
        return jsonify({'error': '没有选择文件'})
    
    # 检查文件类型是否允许
    if file and allowed_file(file.filename):
        # 安全地处理文件名（避免路径遍历等安全问题）
        filename = secure_filename(file.filename)
        # 构建文件保存路径
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        # 保存上传的文件
        file.save(filepath)
        
        try:
            # 加载图像并转换为RGB格式（确保图像格式一致）
            image = Image.open(filepath).convert('RGB')
            # 应用预处理转换并添加批次维度，然后移动到指定设备
            image_tensor = transform(image).unsqueeze(0).to(device)
            
            # 使用模型进行预测（不计算梯度）
            with torch.no_grad():
                # 前向传播获取模型输出
                outputs = model(image_tensor)
                # 获取预测类别（最大概率对应的索引）
                _, predicted = torch.max(outputs, 1)
                
                # 获取预测类别的数字标签
                predicted_class = predicted.item()
                # 将数字标签转换为可读的类别名称
                class_name = class_names.get(predicted_class, "未知")
                
                # 计算预测概率
                probabilities = torch.nn.functional.softmax(outputs, dim=1)[0]
                # 获取预测类别的置信度（百分比）
                confidence = probabilities[predicted_class].item() * 100
                
                # 返回预测结果
                return jsonify({
                    'success': True,
                    'class': class_name,
                    'confidence': f"{confidence:.2f}%",
                    'filename': filename
                })
        
        except Exception as e:
            # 处理预测过程中的错误
            return jsonify({'error': f'预测过程中出错: {str(e)}'})
    
    # 如果文件类型不允许
    return jsonify({'error': '不支持的文件类型'})

if __name__ == '__main__':
    # 当直接运行此文件时，启动Web服务器
    app.run(host='127.0.0.1', port=5000, debug=True) 